import torch

from ai import AiConstant, store_data

# 例是基于Learning PyTorch with Examples这篇教程的，但是我做了一些修改，以适应ctc损失函数和语音数据集。请注意，这个示例只是为了演示ctc
# 损失函数的作用和空白符号的设计，并不是一个完整的语音识别系统。您可能需要根据您的具体需求和数据集进行调整。

# 导入所需的库
import torch
import torchaudio
import numpy as np


# 定义一个简单的CTC模型，使用一个线性层和一个softmax层
class CTCModel(torch.nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super(CTCModel, self).__init__()
        self.linear = torch.nn.Linear(input_size, hidden_size)
        self.softmax = torch.nn.LogSoftmax(dim=2)
        self.output_size = output_size

    def forward(self, x):
        x = self.linear(x)  # (batch_size, seq_length, hidden_size)
        x = self.softmax(x)  # (batch_size, seq_length, output_size)
        return x


# 定义一些超参数
batch_size = 16  # 批次大小
input_size = 13  # 输入特征维度，假设使用MFCC特征
hidden_size = 32  # 隐藏层大小
output_size = 28  # 输出类别数，假设使用26个字母加空白符号和结束符号
learning_rate = 1e-3  # 学习率

# 创建一个随机的语音数据集，每个样本是一个MFCC特征矩阵和一个文本标签
# 这里只是为了演示，实际应用中应该使用真实的语音数据集，如LibriSpeech或Speech Commands
speech_data = []
for i in range(100):
    # 随机生成一个序列长度在10到20之间的MFCC特征矩阵
    seq_length = np.random.randint(10, 20)
    feature = np.random.randn(seq_length, input_size)
    # 随机生成一个长度在5到10之间的文本标签，使用数字0到25表示26个字母
    label_length = np.random.randint(5, 10)
    label = np.random.randint(0, 26, size=label_length)
    # 将特征矩阵和文本标签转换为张量，并添加到数据集中
    feature = torch.from_numpy(feature).float()
    label = torch.from_numpy(label).int()
    speech_data.append((feature, label))

# 定义一个数据加载器，将数据集分成批次，并打乱顺序
speech_loader = torch.utils.data.DataLoader(speech_data, batch_size=batch_size, shuffle=True)

# 创建一个CTC模型实例，并将其移动到设备上（CPU或GPU）
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = CTCModel(input_size, hidden_size, output_size).to(device)

# 创建一个CTC损失函数实例
ctc_loss = torch.nn.CTCLoss(blank=26)  # 假设空白符号的索引是26

# 创建一个优化器实例，使用Adam算法
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)

# 训练模型
num_epochs = 10  # 训练轮数
for epoch in range(num_epochs):
    print(f"Epoch {epoch + 1}")
    for batch_idx, (feature, label) in enumerate(speech_loader):
        # 将特征和标签移动到设备上
        feature = feature.to(device)
        label = label.to(device)
        # 计算每个样本的特征长度和标签长度，这些信息需要传给CTC损失函数
        feature_length = torch.full((batch_size,), feature.size(1), dtype=torch.long)
        label_length = torch.tensor([len(l) for l in label])
        # 前向传播，得到模型输出
        output = model(feature)  # (batch_size, seq_length, output_size)
        # 计算CTC损失函数
        loss = ctc_loss(output, label, feature_length, label_length)
        # 反向传播和参数更新
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        # 打印损失值
        print(f"Batch {batch_idx}, Loss {loss.item()}")
